# -*- coding: utf-8 -*-
import os.path
import numpy as np
import sim_skynet

#%%
fpath = "../data/feature_bin"
feature = np.load("../data/feature_int32_HWC_big.npz")
weight = np.load("../data/skynet_weight_int32_KKC_OCIC.npz")

#%%
def write_bin(img, prefix):
    assert np.all(np.logical_and(-32768 <= img, img < 32768))
    if np.all(np.logical_and(0 <= img, img < 256)):
        dtp = "u8"
        arr = img.astype(np.uint8)
    else:
        dtp = "s16"
        arr = img.astype(np.int16)
    fname = "{}_{}.bin".format(prefix, dtp)
    arr.tofile(os.path.join(fpath, fname))
    return

#%%
val = feature["big"]
write_bin(val, "Raw")

# Resize
val = sim_skynet.zq_resize(val)
sim_skynet.check(val, feature["data0"], "in")
write_bin(val, "Resize")

# Bundle 1
val = sim_skynet.dwconv(val, weight["L0-W"], 128)
write_bin(val, "L0")
val = sim_skynet.norm_act(val, weight["L0-B"], weight["L0-M"])
write_bin(val, "L0_ACTV")
sim_skynet.check(val, feature["conv1"], "L0")

val = sim_skynet.pwconv(val, weight["L1-W"])
write_bin(val, "L1")
val = sim_skynet.norm_act(val, weight["L1-B"], weight["L1-M"])
write_bin(val, "L1_ACTV")
sim_skynet.check(val, feature["conv2"], "L1")

val = sim_skynet.maxPool(val)
sim_skynet.check(val, feature["pool1"], "P1")
write_bin(val, "P1")

# Bundle 2
val = sim_skynet.dwconv(val, weight["L2-W"])
write_bin(val, "L2")
val = sim_skynet.norm_act(val, weight["L2-B"], weight["L2-M"])
write_bin(val, "L2_ACTV")
sim_skynet.check(val, feature["conv3"], "L2")

val = sim_skynet.pwconv(val, weight["L3-W"])
write_bin(val, "L3")
val = sim_skynet.norm_act(val, weight["L3-B"], weight["L3-M"])
write_bin(val, "L3_ACTV")
sim_skynet.check(val, feature["conv4"], "L3")

val = sim_skynet.maxPool(val)
sim_skynet.check(val, feature["pool2"], "P2")
write_bin(val, "P2")

# Bundle 3
val = sim_skynet.dwconv(val, weight["L4-W"])
write_bin(val, "L4")
val = sim_skynet.norm_act(val, weight["L4-B"], weight["L4-M"])
write_bin(val, "L4_ACTV")
sim_skynet.check(val, feature["conv5"], "L4")

val = sim_skynet.pwconv(val, weight["L5-W"])
write_bin(val, "L5")
val = sim_skynet.norm_act(val, weight["L5-B"], weight["L5-M"])
write_bin(val, "L5_ACTV")
sim_skynet.check(val, feature["conv6"], "L5")

# ByPass & ReOrder
sp = val.shape
bypass = val.reshape(sp[0] // 2, 2, sp[1] // 2, 2 * sp[2]).transpose((0, 2, 1, 3))
assert bypass.shape == (sp[0] // 2, sp[1] // 2, 2, 2 * sp[2])
bypass = bypass.reshape(sp[0] // 2, sp[1] // 2, sp[2] * 4)
sim_skynet.check(bypass, feature["reorg"], "Reorg")
write_bin(bypass, "Reorg")

val = sim_skynet.maxPool(val)
sim_skynet.check(val, feature["pool3"], "P3")
write_bin(val, "P3")

# Bundle 4
val = sim_skynet.dwconv(val, weight["L6-W"])
write_bin(val, "L6")
val = sim_skynet.norm_act(val, weight["L6-B"], weight["L6-M"])
write_bin(val, "L6_ACTV")
sim_skynet.check(val, feature["conv7"], "L6")

val = sim_skynet.pwconv(val, weight["L7-W"])
write_bin(val, "L7")
val = sim_skynet.norm_act(val, weight["L7-B"], weight["L7-M"])
write_bin(val, "L7_ACTV")
sim_skynet.check(val, feature["conv8"], "L7")

# Bundle 5
val = sim_skynet.dwconv(val, weight["L8-W"])
write_bin(val, "L8")
val = sim_skynet.norm_act(val, weight["L8-B"], weight["L8-M"])
write_bin(val, "L8_ACTV")
sim_skynet.check(val, feature["conv9"], "L8")

val = sim_skynet.pwconv(val, weight["L9-W"])
write_bin(val, "L9")
val = sim_skynet.norm_act(val, weight["L9-B"], weight["L9-M"])
write_bin(val, "L9_ACTV")
sim_skynet.check(val, feature["conv10"], "L9")

# concatenate
val = np.concatenate((bypass, val), axis=2)
sim_skynet.check(val, feature["cat"], "Concat")
write_bin(val, "Concat")

# Bundle 6
val = sim_skynet.dwconv(val, weight["L10-W"])
write_bin(val, "L10")
val = sim_skynet.norm_act(val, weight["L10-B"], weight["L10-M"])
write_bin(val, "L10_ACTV")
sim_skynet.check(val, feature["conv11"], "L10")

val = sim_skynet.pwconv(val, weight["L11-W"])
write_bin(val, "L11")
val = sim_skynet.norm_act(val, weight["L11-B"], weight["L11-M"])
write_bin(val, "L11_ACTV")
sim_skynet.check(val, feature["conv12"], "L11")

# Output
val = sim_skynet.pwconv(val, weight["L12-W"])
write_bin(val, "L12")
sim_skynet.check(val, feature["conv13"], "L12")

# findmax
val = sim_skynet.find_max(val)
write_bin(val, "Out")
sim_skynet.check(val.reshape(1, *val.shape), feature["fmax"], "out")
